import os
import torch
import pickle
from datetime import datetime
import logging

logger = logging.getLogger(__name__)

class ModelManager:
    def __init__(self, model_dir='models'):
        self.model_dir = model_dir
        os.makedirs(model_dir, exist_ok=True)

    def save_model(self, model, feature_scaler, model_name):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        model_path = os.path.join(self.model_dir, f'{model_name}_{timestamp}.pth')
        scaler_path = os.path.join(self.model_dir, f'{model_name}_scaler_{timestamp}.pkl')
        torch.save(model.state_dict(), model_path)
        with open(scaler_path, 'wb') as f:
            pickle.dump(feature_scaler, f)
        logger.info(f"模型保存: {model_path}")
        return model_path, scaler_path